import os
from PySide6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QPushButton, QFileDialog,
    QTextEdit, QLabel, QLineEdit, QHBoxLayout, QGroupBox,
    QProgressBar, QMessageBox, QSplitter, QFrame
)
from PySide6.QtCore import Qt, QThread, Signal
from PySide6.QtGui import QFont, QPalette, QColor

class ScanThread(QThread):
    """扫描线程，避免界面卡顿"""
    progress_signal = Signal(int)
    text_signal = Signal(str)
    finished_signal = Signal(str)
    
    def __init__(self, root_dir, ignore_extensions):
        super().__init__()
        self.root_dir = root_dir
        self.ignore_extensions = ignore_extensions
        
    def run(self):
        result_text = []
        total_files = 0
        processed_files = 0
        
        # 先统计文件总数
        for dirpath, dirnames, filenames in os.walk(self.root_dir):
            for filename in filenames:
                if not any(filename.endswith(ext) for ext in self.ignore_extensions):
                    total_files += 1
        
        for dirpath, dirnames, filenames in os.walk(self.root_dir):
            filenames.sort()
            for filename in filenames:
                if any(filename.endswith(ext) for ext in self.ignore_extensions):
                    continue

                file_path = os.path.join(dirpath, filename)
                relative_path = os.path.relpath(file_path, self.root_dir)
                
                result_text.append(f"=====FILE_START: {relative_path}=====")
                try:
                    with open(file_path, "r", encoding="utf-8") as f:
                        result_text.append(f.read())
                except Exception as e:
                    result_text.append(f"// 无法读取文件: {e}")
                result_text.append(f"=====FILE_END=====\n")
                
                processed_files += 1
                progress = int((processed_files / total_files) * 100) if total_files > 0 else 0
                self.progress_signal.emit(progress)
        
        final_text = "\n".join(result_text)
        self.finished_signal.emit(final_text)

class FileTool(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("文件扫描与恢复工具 - Professional")
        self.resize(1000, 700)
        
        # 设置现代风格
        self.setStyleSheet("""
            QWidget {
                font-family: 'Segoe UI', 'Microsoft YaHei', sans-serif;
                font-size: 10pt;
            }
            QGroupBox {
                font-weight: bold;
                border: 2px solid #cccccc;
                border-radius: 8px;
                margin-top: 1ex;
                padding-top: 10px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px 0 5px;
            }
            QPushButton {
                background-color: #4CAF50;
                border: none;
                color: white;
                padding: 8px 16px;
                border-radius: 4px;
                font-weight: bold;
                min-width: 80px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
            QPushButton:pressed {
                background-color: #3d8b40;
            }
            QPushButton:disabled {
                background-color: #cccccc;
                color: #666666;
            }
            QTextEdit {
                border: 1px solid #cccccc;
                border-radius: 4px;
                padding: 4px;
            }
            QLineEdit {
                border: 1px solid #cccccc;
                border-radius: 4px;
                padding: 6px;
            }
            QLabel {
                color: #333333;
            }
            QProgressBar {
                border: 1px solid #cccccc;
                border-radius: 4px;
                text-align: center;
            }
            QProgressBar::chunk {
                background-color: #4CAF50;
                border-radius: 3px;
            }
        """)

        # 默认屏蔽类型
        self.ignore_extensions = [".py", ".dll", ".lib"]
        self.scan_thread = None

        self.init_ui()

    def init_ui(self):
        main_layout = QVBoxLayout()
        main_layout.setSpacing(15)
        main_layout.setContentsMargins(20, 20, 20, 20)

        # 标题
        title_label = QLabel("文件扫描与恢复工具")
        title_label.setAlignment(Qt.AlignCenter)
        title_font = QFont()
        title_font.setPointSize(16)
        title_font.setBold(True)
        title_label.setFont(title_font)
        title_label.setStyleSheet("color: #2c3e50; margin-bottom: 10px;")
        main_layout.addWidget(title_label)

        # 设置分组
        settings_group = QGroupBox("基本设置")
        settings_layout = QVBoxLayout()

        # 屏蔽类型设置
        ignore_layout = QVBoxLayout()
        ignore_layout.addWidget(QLabel("屏蔽的文件类型:"))
        self.ignore_input = QLineEdit(",".join(self.ignore_extensions))
        self.ignore_input.setPlaceholderText("输入要屏蔽的文件扩展名，用逗号分隔，如：.py,.dll,.lib")
        ignore_layout.addWidget(self.ignore_input)
        settings_layout.addLayout(ignore_layout)

        settings_group.setLayout(settings_layout)
        main_layout.addWidget(settings_group)

        # 扫描分组
        scan_group = QGroupBox("文件扫描")
        scan_layout = QVBoxLayout()

        # 扫描目录选择
        scan_dir_layout = QHBoxLayout()
        scan_dir_layout.addWidget(QLabel("扫描目录:"))
        self.scan_dir_input = QLineEdit()
        self.scan_dir_input.setPlaceholderText("请选择要扫描的目录...")
        scan_dir_layout.addWidget(self.scan_dir_input)
        self.scan_btn = QPushButton("浏览")
        self.scan_btn.clicked.connect(self.choose_scan_dir)
        scan_dir_layout.addWidget(self.scan_btn)
        scan_layout.addLayout(scan_dir_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        scan_layout.addWidget(self.progress_bar)

        # 扫描按钮
        self.scan_generate_btn = QPushButton("开始扫描并生成文本")
        self.scan_generate_btn.clicked.connect(self.scan_and_generate)
        self.scan_generate_btn.setStyleSheet("background-color: #2196F3;")
        scan_layout.addWidget(self.scan_generate_btn)

        scan_group.setLayout(scan_layout)
        main_layout.addWidget(scan_group)

        # 文本显示区域
        text_group = QGroupBox("文件内容预览")
        text_layout = QVBoxLayout()
        self.scan_text = QTextEdit()
        self.scan_text.setPlaceholderText("扫描结果将显示在这里...")
        text_layout.addWidget(self.scan_text)
        text_group.setLayout(text_layout)
        main_layout.addWidget(text_group)

        # 恢复分组
        restore_group = QGroupBox("文件恢复")
        restore_layout = QVBoxLayout()

        # 文本文件选择
        restore_file_layout = QHBoxLayout()
        restore_file_layout.addWidget(QLabel("文本文件:"))
        self.restore_file_input = QLineEdit()
        self.restore_file_input.setPlaceholderText("请选择包含文件数据的文本文件...")
        restore_file_layout.addWidget(self.restore_file_input)
        self.restore_btn = QPushButton("浏览")
        self.restore_btn.clicked.connect(self.choose_restore_file)
        restore_file_layout.addWidget(self.restore_btn)
        restore_layout.addLayout(restore_file_layout)

        # 输出目录
        output_layout = QHBoxLayout()
        output_layout.addWidget(QLabel("输出目录:"))
        self.output_dir_input = QLineEdit("recovered_files")
        output_layout.addWidget(self.output_dir_input)
        restore_layout.addLayout(output_layout)

        # 恢复按钮
        self.restore_generate_btn = QPushButton("从文本生成文件")
        self.restore_generate_btn.clicked.connect(self.restore_files)
        self.restore_generate_btn.setStyleSheet("background-color: #FF9800;")
        restore_layout.addWidget(self.restore_generate_btn)

        restore_group.setLayout(restore_layout)
        main_layout.addWidget(restore_group)

        # 状态栏
        self.status_label = QLabel("就绪")
        self.status_label.setStyleSheet("color: #666; font-size: 9pt; padding: 5px;")
        main_layout.addWidget(self.status_label)

        self.setLayout(main_layout)

    def choose_scan_dir(self):
        folder = QFileDialog.getExistingDirectory(self, "选择扫描目录")
        if folder:
            self.scan_dir_input.setText(folder)
            self.status_label.setText(f"已选择目录: {folder}")

    def scan_and_generate(self):
        root_dir = self.scan_dir_input.text()
        if not root_dir or not os.path.exists(root_dir):
            QMessageBox.warning(self, "警告", "请选择有效的扫描目录！")
            return

        # 动态更新屏蔽类型
        self.ignore_extensions = [ext.strip() for ext in self.ignore_input.text().split(",") if ext.strip()]

        # 禁用按钮，显示进度条
        self.scan_generate_btn.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.status_label.setText("正在扫描文件，请稍候...")

        # 创建并启动扫描线程
        self.scan_thread = ScanThread(root_dir, self.ignore_extensions)
        self.scan_thread.progress_signal.connect(self.update_progress)
        self.scan_thread.text_signal.connect(self.scan_text.append)
        self.scan_thread.finished_signal.connect(self.scan_finished)
        self.scan_thread.start()

    def update_progress(self, value):
        self.progress_bar.setValue(value)

    def scan_finished(self, final_text):
        self.scan_text.setPlainText(final_text)
        
        # 保存文件
        output_file = "all_files_combined.txt"
        try:
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(final_text)
            self.status_label.setText(f"扫描完成！结果已保存到: {output_file}")
        except Exception as e:
            self.status_label.setText(f"保存文件时出错: {str(e)}")
        
        # 恢复UI状态
        self.scan_generate_btn.setEnabled(True)
        self.progress_bar.setVisible(False)
        QMessageBox.information(self, "完成", f"文件扫描完成！\n共处理了 {final_text.count('=====FILE_START:')} 个文件。")

    def choose_restore_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择文本文件", "", 
            "Text Files (*.txt);;All Files (*)"
        )
        if file_path:
            self.restore_file_input.setText(file_path)
            self.status_label.setText(f"已选择文件: {os.path.basename(file_path)}")

    def restore_files(self):
        input_file = self.restore_file_input.text()
        output_root = self.output_dir_input.text() or "recovered_files"
        
        if not input_file or not os.path.exists(input_file):
            QMessageBox.warning(self, "警告", "请选择有效的文本文件！")
            return

        try:
            # 显示进度
            self.restore_generate_btn.setEnabled(False)
            self.status_label.setText("正在恢复文件...")

            files_restored = self.restore_files_from_text(input_file, output_root)
            
            self.status_label.setText(f"恢复完成！共恢复了 {files_restored} 个文件到: {output_root}")
            QMessageBox.information(self, "完成", 
                                  f"文件恢复完成！\n共恢复了 {files_restored} 个文件到目录:\n{output_root}")
            
        except Exception as e:
            self.status_label.setText(f"恢复文件时出错: {str(e)}")
            QMessageBox.critical(self, "错误", f"恢复文件时出现错误:\n{str(e)}")
        finally:
            self.restore_generate_btn.setEnabled(True)

    def restore_files_from_text(self, input_file, output_root):
        if not os.path.exists(output_root):
            os.makedirs(output_root)

        current_file = None
        content_lines = []
        files_restored = 0

        with open(input_file, "r", encoding="utf-8") as f:
            for line in f:
                line_strip = line.strip()
                if line_strip.startswith("=====FILE_START:"):
                    if current_file:
                        self.write_file(current_file, content_lines)
                        files_restored += 1
                        content_lines = []
                    relative_path = line_strip[len("=====FILE_START:"):].replace("=====", "").strip()
                    current_file = os.path.join(output_root, relative_path)
                elif line_strip == "=====FILE_END=====":
                    if current_file:
                        self.write_file(current_file, content_lines)
                        files_restored += 1
                        current_file = None
                        content_lines = []
                else:
                    if current_file is not None:
                        content_lines.append(line)

        # 处理最后一个文件
        if current_file and content_lines:
            self.write_file(current_file, content_lines)
            files_restored += 1

        return files_restored

    def write_file(self, file_path, lines):
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.writelines(lines)


if __name__ == "__main__":
    app = QApplication([])
    
    # 设置应用程序样式
    app.setStyle('Fusion')
    
    window = FileTool()
    window.show()
    app.exec()